#include <cuda_runtime.h>
#include <cusparse.h>
#include <stdio.h>
#include <stdlib.h>

const int Nx = 192;
const int Ny = 256;
const double dr = 0.010471;
const double dz = 0.012548;

void loadfield(double field[Nx*Ny], const char* str)
{
    int i = 0;
    int j = 0;
    FILE *fp = NULL;
    char *line,*record;
    char buffer[20450];
    if((fp = fopen(str, "r")) != NULL)
    {
        //printf("file open success:%s\n",str);
        while ((line = fgets(buffer, sizeof(buffer), fp))!=NULL)
        {
            record = strtok(line, ",");
            while (record != NULL)  
            {
                field[i*Ny+j++] = atof(record);
                record = strtok(NULL, ",");
            }
            i++; j = 0;
        }
        fclose(fp);
        fp = NULL;
        //printf("file close success:%s\n",str);
    }
}

void cudaMv(int* offsets, int* columns, double* values, double* vec, int mdim, int nval, double* retar)
{
    int *dcsr_offsets, *dcsr_columns;
    double* dcsr_values, *dpsi, *ddifx; 
    double alpha = 1.0f;
    double beta = 0.0f;

    cudaMalloc((void**)&dcsr_offsets, (mdim+1)*sizeof(int));
    cudaMalloc((void**)&dcsr_columns, nval*sizeof(int));
    cudaMalloc((void**)&dcsr_values, nval*sizeof(double));
    cudaMalloc((void**)&dpsi, mdim*sizeof(double));
    cudaMalloc((void**)&ddifx, mdim*sizeof(double));
    cudaMemcpy(dcsr_offsets, offsets, (mdim+1)*sizeof(int),cudaMemcpyHostToDevice);
    cudaMemcpy(dcsr_columns, columns, nval*sizeof(int),cudaMemcpyHostToDevice);
    cudaMemcpy(dcsr_values, values, nval*sizeof(double),cudaMemcpyHostToDevice);
    cudaMemcpy(dpsi, vec, mdim*sizeof(double),cudaMemcpyHostToDevice);

    cusparseHandle_t handle;
    cusparseSpMatDescr_t mat;
    cusparseDnVecDescr_t vecpsi;
    cusparseDnVecDescr_t vecdifx;
    void* dbuf = NULL;
    size_t bufsize = 0;

    cusparseCreate(&handle);
    cusparseCreateCsr(&mat, mdim, mdim, nval, dcsr_offsets, dcsr_columns,dcsr_values,
    CUSPARSE_INDEX_32I, CUSPARSE_INDEX_32I, CUSPARSE_INDEX_BASE_ZERO, CUDA_R_64F); 
    cusparseCreateDnVec(&vecpsi, mdim, dpsi, CUDA_R_64F);
    cusparseCreateDnVec(&vecdifx, mdim, ddifx, CUDA_R_64F);

    cusparseSpMV_bufferSize(handle, CUSPARSE_OPERATION_NON_TRANSPOSE,
    &alpha, mat, vecpsi, &beta, vecdifx, CUDA_R_64F, CUSPARSE_SPMV_ALG_DEFAULT, &bufsize);    
    cudaMalloc(&dbuf, bufsize);
    cusparseSpMV(handle, CUSPARSE_OPERATION_NON_TRANSPOSE,
    &alpha, mat, vecpsi, &beta, vecdifx, CUDA_R_64F, CUSPARSE_SPMV_ALG_DEFAULT, dbuf);

    cudaMemcpy(retar, ddifx, mdim*sizeof(double),cudaMemcpyDeviceToHost);
    return;
}

void matasem_dy(int* offsets, int* columns, double* vals, int mdim)
{
    for(int i = 0; i < mdim; i++)
    {
        if(i == 0)
            offsets[i] = 0;
        offsets[i+1] = offsets[i] + 2;
        int ix = i/Ny;
        int iy = i%Ny;

        if (ix == 0)
        {
            columns[2*i] = i;
            vals[2*i] = -1.0/dz;
            columns[2*i+1] = i+Ny;
            vals[2*i+1] = 1.0/dz;
        } else if (ix == Nx-1)
        {
            columns[2*i] = i-Ny;
            vals[2*i] = -1.0/dz;
            columns[2*i+1] = i;
            vals[2*i+1] = 1.0/dz;
        } else
        {
            columns[2*i] = i-Ny;
            columns[2*i+1] = i+Ny;
            vals[2*i] = -1.0/(2.0*dz);
            vals[2*i+1] = 1.0/(2.0*dz);
        }
    }
    return;
}

// not verified
void matasem_dx(int* offsets, int* columns, double* vals, int mdim)
{
    for(int i = 0; i < mdim; i++)
    {
        if(i == 0)
            offsets[i] = 0;
        offsets[i+1] = offsets[i] + 2;
        int ix = i/Ny;
        int iy = i%Ny;

        if (iy == 0)
        {
            columns[2*i] = i;
            vals[2*i] = -1.0/dr;
            columns[2*i+1] = i+1;
            vals[2*i+1] = 1.0/dr;
        } else if (iy == Ny-1)
        {
            columns[2*i] = i-1;
            vals[2*i] = -1.0/dr;
            columns[2*i+1] = i;
            vals[2*i+1] = 1.0/dr;
        } else
        {
            columns[2*i] = i-1;
            columns[2*i+1] = i+1;
            vals[2*i] = -1.0/(2.0*dr);
            vals[2*i+1] = 1.0/(2.0*dr);
        }
    }
    return;
}

int main(void)
{
    double psi[Nx*Ny];
    loadfield(psi,"psi.csv");   
    int nval = 2*Nx*Ny;
    int csr_offsets[Nx*Ny+1] = {0};    
    int csr_columns[nval] = {0};
    double csr_values[nval] = {0.0};
    double difx[Nx*Ny] = {0.0};

    matasem_dy(csr_offsets, csr_columns, csr_values, Nx*Ny);
    
    double ret[Nx*Ny] = {0.0};
    cudaMv(csr_offsets, csr_columns,csr_values,psi, Nx*Ny, 2*Nx*Ny, ret);
    for(int i = 256*54; i < 256*54+10; i++)
        printf("%lf, %lf\n",psi[i], *(ret+i));
    
    return 0;
}